import json
from typing import Dict
from pathlib import Path
from src.huffman import HuffmanTree

class HuffmanExtractor:
    def __init__(self, input_path: Path, output_path: Path):
        """哈夫曼树解压器初始化函数

        Args:
            input_path (Path): 输入文件
            output_path (Path): 输出文件
        """
        self.input_path = input_path
        self.output_path = output_path
        self.tree = HuffmanTree()

    def _read_metadata(self, file) -> Dict[str, int]:
        """读取压缩文件中的元数据（字频字典）

        Args:
            file: 已打开的文件对象

        Returns:
            Dict[str, int]: 解压需要的字频字典
        """
        metadata_len = int.from_bytes(file.read(4), 'big')
        metadata = file.read(metadata_len)
        return json.loads(metadata.decode('utf-8'))

    def _to_bit_string(self, byte_data: bytes) -> str:
        """将字节序列转换为二进制字符串

        Args:
            byte_data (bytes): 字节序列

        Returns:
            str: 转换后的二进制字符串
        """
        padding = byte_data[0]
        bit_string = ''.join(f'{byte:08b}' for byte in byte_data[1:])
        return bit_string[:-padding] if padding else bit_string

    def _decode(self, bit_string: str, codes: Dict[str, str]) -> str:
        """根据哈夫曼编码解码二进制字符串

        Args:
            bit_string (str): 二进制字符串
            codes (Dict[str, str]): 哈夫曼编码字典

        Returns:
            str: 解码后的原始内容
        """
        reverse_codes = {v: k for k, v in codes.items()}
        current_code = ""
        decoded_content = []

        for bit in bit_string:
            current_code += bit
            if current_code in reverse_codes:
                decoded_content.append(reverse_codes[current_code])
                current_code = ""
        
        return ''.join(decoded_content)

    def extract(self) -> None:
        if not self.input_path.exists():
            print(f"输入文件不存在：{self.input_path}")
            return

        with self.input_path.open('rb') as file:
            freq_dict = self._read_metadata(file)
            self.tree.build_tree(freq_dict)
            codes = self.tree.get_codes()

            byte_data = file.read()
            bit_string = self._to_bit_string(byte_data)

            decoded_content = self._decode(bit_string, codes)

        with self.output_path.open('w', encoding='utf-8') as file:
            file.write(decoded_content)

        print(f"解压完成，保存至：{self.output_path}")

if __name__ == "__main__":
    input_path = Path("compressed_HP.txt")
    output_path = Path("origin_HP.txt")
    extractor = HuffmanExtractor(input_path, output_path)
    extractor.extract()